package balancer

import (
	"hash/fnv"
	"math/rand"
	"sync"
	"sync/atomic"
	"time"
)

// Backend 表示一个后端服务器
type Backend struct {
	URL      string
	Weight   int
	Health   bool
	Failures int64
}

// Balancer 定义负载均衡器接口
type Balancer interface {
	Next() *Backend
	Update([]*Backend)
	MarkFailed(*Backend)
	MarkAlive(*Backend)
	Close() // 添加关闭方法
}

// RoundRobinBalancer 实现轮询算法
type RoundRobinBalancer struct {
	backends []*Backend
	current  uint32
	mu       sync.RWMutex
}

func NewRoundRobinBalancer(backends []*Backend) *RoundRobinBalancer {
	return &RoundRobinBalancer{
		backends: backends,
	}
}

func (b *RoundRobinBalancer) Next() *Backend {
	b.mu.RLock()
	defer b.mu.RUnlock()

	if len(b.backends) == 0 {
		return nil
	}

	next := atomic.AddUint32(&b.current, 1)
	return b.backends[next%uint32(len(b.backends))]
}

func (b *RoundRobinBalancer) Update(backends []*Backend) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = backends
}

func (b *RoundRobinBalancer) MarkFailed(backend *Backend) {
	atomic.AddInt64(&backend.Failures, 1)
	backend.Health = false
}

func (b *RoundRobinBalancer) MarkAlive(backend *Backend) {
	atomic.StoreInt64(&backend.Failures, 0)
	backend.Health = true
}

// WeightedRoundRobinBalancer 实现加权轮询算法
type WeightedRoundRobinBalancer struct {
	backends       []*Backend
	currentWeights []int
	mu             sync.RWMutex
}

func NewWeightedRoundRobinBalancer(backends []*Backend) *WeightedRoundRobinBalancer {
	weights := make([]int, len(backends))
	return &WeightedRoundRobinBalancer{
		backends:       backends,
		currentWeights: weights,
	}
}

func (b *WeightedRoundRobinBalancer) Next() *Backend {
	b.mu.Lock()
	defer b.mu.Unlock()

	if len(b.backends) == 0 {
		return nil
	}

	totalWeight := 0
	maxWeight := -1
	selectedIndex := -1

	for i, backend := range b.backends {
		if !backend.Health {
			continue
		}
		b.currentWeights[i] += backend.Weight
		totalWeight += backend.Weight

		if maxWeight < b.currentWeights[i] {
			maxWeight = b.currentWeights[i]
			selectedIndex = i
		}
	}

	if selectedIndex >= 0 {
		b.currentWeights[selectedIndex] -= totalWeight
		return b.backends[selectedIndex]
	}

	return nil
}

func (b *WeightedRoundRobinBalancer) Update(backends []*Backend) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = backends
	b.currentWeights = make([]int, len(backends))
}

func (b *WeightedRoundRobinBalancer) MarkFailed(backend *Backend) {
	atomic.AddInt64(&backend.Failures, 1)
	backend.Health = false
}

func (b *WeightedRoundRobinBalancer) MarkAlive(backend *Backend) {
	atomic.StoreInt64(&backend.Failures, 0)
	backend.Health = true
}

// ConsistentHashBalancer 实现一致性哈希算法
type ConsistentHashBalancer struct {
	backends []*Backend
	mu       sync.RWMutex
}

func NewConsistentHashBalancer(backends []*Backend) *ConsistentHashBalancer {
	return &ConsistentHashBalancer{
		backends: backends,
	}
}

func (b *ConsistentHashBalancer) hash(key string) uint32 {
	h := fnv.New32a()
	h.Write([]byte(key))
	return h.Sum32()
}

func (b *ConsistentHashBalancer) NextWithKey(key string) *Backend {
	b.mu.RLock()
	defer b.mu.RUnlock()

	if len(b.backends) == 0 {
		return nil
	}

	hash := b.hash(key)
	index := hash % uint32(len(b.backends))
	return b.backends[index]
}

func (b *ConsistentHashBalancer) Next() *Backend {
	// 如果没有提供key，使用随机选择
	return b.NextWithKey(time.Now().String())
}

func (b *ConsistentHashBalancer) Update(backends []*Backend) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = backends
}

func (b *ConsistentHashBalancer) MarkFailed(backend *Backend) {
	atomic.AddInt64(&backend.Failures, 1)
	backend.Health = false
}

func (b *ConsistentHashBalancer) MarkAlive(backend *Backend) {
	atomic.StoreInt64(&backend.Failures, 0)
	backend.Health = true
}

// RandomBalancer 实现随机算法
type RandomBalancer struct {
	backends []*Backend
	mu       sync.RWMutex
	rand     *rand.Rand
}

func NewRandomBalancer(backends []*Backend) *RandomBalancer {
	return &RandomBalancer{
		backends: backends,
		rand:     rand.New(rand.NewSource(time.Now().UnixNano())),
	}
}

func (b *RandomBalancer) Next() *Backend {
	b.mu.RLock()
	defer b.mu.RUnlock()

	if len(b.backends) == 0 {
		return nil
	}

	healthyBackends := make([]*Backend, 0)
	for _, backend := range b.backends {
		if backend.Health {
			healthyBackends = append(healthyBackends, backend)
		}
	}

	if len(healthyBackends) == 0 {
		return nil
	}

	return healthyBackends[b.rand.Intn(len(healthyBackends))]
}

func (b *RandomBalancer) Update(backends []*Backend) {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = backends
}

func (b *RandomBalancer) MarkFailed(backend *Backend) {
	atomic.AddInt64(&backend.Failures, 1)
	backend.Health = false
}

func (b *RandomBalancer) MarkAlive(backend *Backend) {
	atomic.StoreInt64(&backend.Failures, 0)
	backend.Health = true
}

// LeastConnectionBalancer 实现最小连接数算法
type LeastConnectionBalancer struct {
	backends    []*Backend
	connections map[string]*int64
	mu          sync.RWMutex
}

func NewLeastConnectionBalancer(backends []*Backend) *LeastConnectionBalancer {
	connections := make(map[string]*int64)
	for _, backend := range backends {
		var count int64
		connections[backend.URL] = &count
	}

	return &LeastConnectionBalancer{
		backends:    backends,
		connections: connections,
	}
}

func (b *LeastConnectionBalancer) Next() *Backend {
	b.mu.RLock()
	defer b.mu.RUnlock()

	if len(b.backends) == 0 {
		return nil
	}

	var selected *Backend
	minConn := int64(^uint64(0) >> 1) // max int64

	for _, backend := range b.backends {
		if !backend.Health {
			continue
		}

		conns := atomic.LoadInt64(b.connections[backend.URL])
		if conns < minConn {
			minConn = conns
			selected = backend
		}
	}

	if selected != nil {
		b.incrementConnection(selected.URL)
	}

	return selected
}

func (b *LeastConnectionBalancer) Update(backends []*Backend) {
	b.mu.Lock()
	defer b.mu.Unlock()

	newConnections := make(map[string]*int64)
	for _, backend := range backends {
		if count, exists := b.connections[backend.URL]; exists {
			newConnections[backend.URL] = count
		} else {
			var count int64
			newConnections[backend.URL] = &count
		}
	}

	b.backends = backends
	b.connections = newConnections
}

func (b *LeastConnectionBalancer) MarkFailed(backend *Backend) {
	atomic.AddInt64(&backend.Failures, 1)
	backend.Health = false
}

func (b *LeastConnectionBalancer) MarkAlive(backend *Backend) {
	atomic.StoreInt64(&backend.Failures, 0)
	backend.Health = true
}

func (b *LeastConnectionBalancer) incrementConnection(backend string) {
	if _, exists := b.connections[backend]; !exists {
		var count int64 = 1
		b.connections[backend] = &count
	} else {
		*b.connections[backend]++
	}
}

func (b *LeastConnectionBalancer) Close() {
	b.mu.Lock()
	defer b.mu.Unlock()

	for k := range b.connections {
		delete(b.connections, k)
	}
	b.backends = nil
}

func (b *RoundRobinBalancer) Close() {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = nil
}

func (b *WeightedRoundRobinBalancer) Close() {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = nil
	b.currentWeights = nil
}

func (b *ConsistentHashBalancer) Close() {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = nil
}

func (b *RandomBalancer) Close() {
	b.mu.Lock()
	defer b.mu.Unlock()
	b.backends = nil
}
